{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'ray'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-2b3eb2205af9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mast\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mliteral_eval\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mray\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mgym\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mretro\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'ray'"
     ]
    }
   ],
   "source": [
    "from ast import literal_eval\n",
    "import ray\n",
    "import gym\n",
    "import retro\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from markov import sampleMarkov, createMarkov, randMarkov\n",
    "from support import getInitial, verifyTrajectory, install_games_from_rom_dir, frameToCell, action_set, trajectoryToGif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import imageio\n",
    "imageio.plugins.freeimage.download()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class GoExploreActor(object):\n",
    "    def __init__(self, name):\n",
    "        self.env = retro.make(name)\n",
    "        self.env.reset()\n",
    "        \n",
    "    def setPolicy(self, policy):\n",
    "        self.policy_type = policy['type']\n",
    "        if self.policy_type=='markov':\n",
    "            self.policy_weights = policy['weights']\n",
    "    \n",
    "    def updateCache(self, cells, fitnesses):\n",
    "        self.cells = cells\n",
    "        self.fitnesses = fitnesses\n",
    "    \n",
    "    def GoExplore(self, state, trajectory, steps):\n",
    "        results = []\n",
    "        visits = []\n",
    "        self.env.em.set_state(state)      \n",
    "        \n",
    "        recurrent_state = None\n",
    "        \n",
    "        if self.policy_type=='markov':\n",
    "            recurrent_state = np.random.randint(12)\n",
    "        \n",
    "        for i in range(steps):\n",
    "            action = None\n",
    "            if self.policy_type == 'random':\n",
    "                action = np.random.randint(12)\n",
    "            if self.policy_type == 'markov':\n",
    "                action = sampleMarkov(recurrent_state, self.policy_weights)\n",
    "                recurrent_state = action\n",
    "                \n",
    "            observation, reward, done, info = self.env.step(action_set[action])\n",
    "            trajectory.append(action)\n",
    "            \n",
    "            cell = frameToCell(observation, info)\n",
    "            \n",
    "            if cell not in visits:\n",
    "                visits.append(cell)\n",
    "                \n",
    "            if cell in self.cells:\n",
    "                if len(trajectory) < self.fitnesses[cell]:\n",
    "                    results.append((cell, trajectory.copy(), self.env.em.get_state(), info, i))\n",
    "            else:\n",
    "                results.append((cell, trajectory.copy(), self.env.em.get_state(), info, i))\n",
    "                self.cells.append(cell)\n",
    "                self.fitnesses[cell] = len(trajectory)\n",
    "        return (results, visits)\n",
    "\n",
    "    def reset(self):\n",
    "        self.env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "install_games_from_rom_dir('roms/')\n",
    "\n",
    "game = 'SonicTheHedgehog-Genesis'\n",
    "stateStr = 'GreenHillZone.Act1.state'\n",
    "\n",
    "policy = {'type':'markov', 'weights':randMarkov(10,12)}\n",
    "\n",
    "initialCell, initialState, initialTrajectory, initialFitness = getInitial(game, stateStr)\n",
    "\n",
    "cells = [initialCell]\n",
    "fitnesses = {initialCell:initialFitness}\n",
    "cell_prob = {initialCell:1}\n",
    "trajectories = {initialCell:initialTrajectory}\n",
    "states = {initialCell:initialState}\n",
    "\n",
    "NWorkers = 2\n",
    "\n",
    "workers = [ GoExploreActor.remote(game) for _ in range(NWorkers)]\n",
    "\n",
    "for worker in workers:\n",
    "    worker.setPolicy.remote(policy)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def winCondition(cell, info):\n",
    "    return info['level_end_bonus'] != 0\n",
    "\n",
    "best_trajectory = None\n",
    "verbose = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stat_cells = np.zeros(500)\n",
    "stat_avglen = np.zeros(500)\n",
    "stat_runlen = np.zeros(500)\n",
    "#Optionally add other numbers of steps to this list to alternate shorter episodes\n",
    "nsteps = [500]\n",
    "\n",
    "for i in range(500):\n",
    "    print('Iteration ' + str(i))\n",
    "    \n",
    "    #Pass the full cell table and fitnesses table to each worker\n",
    "    for worker in workers:\n",
    "        worker.updateCache.remote(ray.put(cells), ray.put(fitnesses))\n",
    "    normalized_cell_prob = np.array([cell_prob[c] for c in cells])\n",
    "    normalized_cell_prob = normalized_cell_prob/normalized_cell_prob.sum()\n",
    "    goCells = [cells[r] for r in np.random.choice(np.arange(len(cells)),size=(NWorkers) , p = normalized_cell_prob )]\n",
    "    #Select a cell randomly\n",
    "    #goCells = [cells[np.random.randint(len(cells))] for i in range(NWorkers)]\n",
    "    R = ray.get( [ workers[i].GoExplore.remote(ray.put(states[goCells[i]]), ray.put(trajectories[goCells[i]]), nsteps[i%len(nsteps)]) for i in range(NWorkers)]   )\n",
    "    \n",
    "    #Complie results and update master tables\n",
    "    results = []\n",
    "    visits = []\n",
    "    for r, v in R:\n",
    "        results+=r\n",
    "        visits+=v\n",
    "    \n",
    "    for cell, trajectory, state, info, iteration in results:\n",
    "        if cell in cells:\n",
    "            if len(trajectory) < fitnesses[cell]:\n",
    "                if verbose:\n",
    "                    print('Shortened trajectory to ' + cell + ' from ' \n",
    "                          + str(fitnesses[cell]) + ' to ' \n",
    "                          + str(len(trajectory)) + ', saving ' \n",
    "                          + str(fitnesses[cell]-len(trajectory)) + ' frames at step '\n",
    "                          + str(iteration))\n",
    "                fitnesses[cell] = len(trajectory)\n",
    "                trajectories[cell] = trajectory\n",
    "                states[cell] = state\n",
    "                cell_prob[cell] += 1\n",
    "        else:\n",
    "            if winCondition(cell, info):\n",
    "                if best_trajectory is None:\n",
    "                    best_trajectory = trajectory\n",
    "                    print('Improved trajectory: ' + str(len(best_trajectory)) + ' Frames')\n",
    "                elif len(trajectory)<len(best_trajectory):\n",
    "                    best_trajectory = trajectory       \n",
    "                    print('Improved trajectory: ' + str(len(best_trajectory)) + ' Frames')             \n",
    "            else:\n",
    "                if verbose:\n",
    "                    print('Found cell ' + cell + ' at step ' + str(iteration) + ', total ' + str(len(cells)) + ' found')\n",
    "                cells.append(cell)\n",
    "                fitnesses[cell] = len(trajectory)\n",
    "                trajectories[cell] = trajectory\n",
    "                states[cell] = state\n",
    "                cell_prob[cell] = 10\n",
    "    \n",
    "    for cell in visits:\n",
    "        cell_prob[cell] *= .5\n",
    "    if best_trajectory is None:\n",
    "        if i < 100 and i%2==0:\n",
    "            policy['weights'] = randMarkov(10,12)\n",
    "        else:\n",
    "            policy['weights'] = createMarkov(trajectories[cells[-1]],12)\n",
    "    else:\n",
    "        policy['weights'] = createMarkov(best_trajectory)\n",
    "        stat_runlen[i] = len(best_trajectory)\n",
    "    \n",
    "    for worker in workers:\n",
    "        worker.setPolicy.remote(policy)\n",
    "        \n",
    "    stat_cells[i] = len(cells)\n",
    "    stat_avglen[i] = np.array(list(fitnesses.values())).mean()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(stat_cells)\n",
    "plt.show()\n",
    "plt.plot(stat_avglen)\n",
    "plt.show()\n",
    "plt.plot(stat_runlen)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "cell = cells[np.array([literal_eval(cell)[0] for cell in cells]).argsort()[-1]]\n",
    "\n",
    "c = literal_eval(cell)\n",
    "trajectory = trajectories[cell]\n",
    "trajectoryToGif(game, stateStr, trajectory, True, 'Gameplay_'+str(len(trajectory))+'-'+str(c[0])+'-'+str(c[1])+'-'+str(c[2])+'-'+str(c[3])+'.gif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
